num_class = 4
input_size = 4
hidden_size = 8
embedding_dim = 10 # 嵌入10维空间
batch_size = 1
num_layers = 1
seq_len = 5
# 准备数据
idx2char = ['e','h','l','o'] # 字典
x_data = [[1,0,2,2,3]]
y_data = [3,1,2,3,2]  # ohlol
inputs = torch.LongTensor(x_data)
labels = torch.LongTensor(y_data)
# 构造模型
class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.emb = nn.Embedding(input_size,embedding_dim)
        #返回inputs_len*embedding_dim
        self.lstm = nn.LSTM(embedding_dim,  # 10
                                 hidden_size,    # 8
                                 num_layers,    # 1
                                 batch_first=True
        # (seq_len,batch_size,input_size)转化成(batch_size,seq_len,input_size)
                                  )
        self.fc = nn.Linear(hidden_size,num_class)

    def forward(self,x):
        x = self.emb(x)
        # (batch,seqlen,embeddingsize) (1,5,10)
        output,(h_n, c_n) = self.lstm(x)
        # lstm(x,(h_0, c_0))如果未提供（h_0,c_0）则默认为零。
        # output:(batch_size,seq_len, input_size) (1,5,8)
        # h_n,c_n: (batch_size,num_layers,hidden_size)
        # h_n:(1,1,8)  c_n:(1,1,8)
        x = self.fc(output)
        # x:(batch_size,seq_len,num_class) (1,5,4)
        return x
model = Model()
# 损失函数和优化器
criterion = nn.CrossEntropyLoss(reduction='mean')
optimizer = torch.optim.Adam(model.parameters(),lr=0.06)
# 训练
for epoch in range(30):
    outputs = model(inputs)
    # inputs是（batch,input_size）(1,5)
    # outputs是(batch,seq,num_class) (1,5,4)

    pred = outputs.permute(0, 2, 1)
    # pred:(batch,num_class,seqlen) (1,4,5)

    labels_y = labels.view(-1, 5)
    # labels_y:（batch,seqlen） (1,5)

    loss = criterion(pred,labels_y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    _,idx = pred.max(dim=1)
    # pred:(batch,num_class,seqlen) (1,4,5)
    # max(dim=1),找到dim=1维度的最大值,
    # 即:消除形状(1,4,5)中dim=1维度,得到idx:(1,5)
    idx = idx.view(5, ).data.numpy()# 将tensor转成numpy数组
    print("Predicted:",''.join([idx2char[x] for x in idx]),end='')
    print(",Epoch {}/30 loss={:.3f}". format(epoch+1,loss.item()) 
